# %%
from collections import defaultdict
from typing import Literal
from conn import (
    mysql_engine,
    clip_annotated_record,
    gt_task_item_col,
    gt_task_col,
    check_frame_label_col,
    ac_task_col,
)
from sqlalchemy import text
from pydantic import BaseModel, Field
import pandas as pd

from enums import AnnotateTypeEnum


GTSupplierCnt = dict[Literal["sensetime", "imotion", "zdrive"], int]


class GTDashboardModel(BaseModel):
    annotate_type: int = Field(default=0, title="标注类型")
    annotate_name: str = Field(default="", title="标注类型名称")

    curated_clip_cnt: int = Field(default=0, title="送标 clip 数")
    curated_frame_cnt: int = Field(default=0, title="送标 frame 数")

    labeled_clip_cnt: int = Field(default=0, title="已标注 clip 数")
    labeled_frame_cnt: int = Field(default=0, title="已标注 frame 数")

    labeled_vis_processed_clip_cnt: int = Field(
        default=0, title="已标注 clip 数(可见性擦除)"
    )
    labeled_vis_processed_frame_cnt: int = Field(
        default=0, title="已标注 frame 数(可见性擦除)"
    )

    accept_pending_clip_cnt: int = Field(default=0, title="待验收 clip 数")
    accept_pending_frame_cnt: int = Field(default=0, title="待验收 frame 数")

    accept_finished_clip_cnt: int = Field(default=0, title="已验收 clip 数")
    accept_finished_frame_cnt: int = Field(default=0, title="已验收 frame 数")
    accept_valid_frame_cnt: int = Field(default=0, title="验收有效 frame 数")
    accept_valid_rate: float = Field(default=0.0, title="验收合格率")

    gt_task_running_clip_cnt: int = Field(default=0, title="真值生产中 clip 数")
    gt_task_running_frame_cnt: int = Field(default=0, title="真值生产中 frame 数")

    gt_task_success_clip_cnt: int = Field(default=0, title="真值生产完成 clip 数")
    gt_task_success_frame_cnt: int = Field(default=0, title="真值生产完成 frame 数")

    gt_task_running_clip_cnt_by_supplier: GTSupplierCnt = Field(
        default_factory=lambda: defaultdict(int),
        title="不同供应商的真值生产中 clip 数",
    )
    gt_task_running_frame_cnt_by_supplier: GTSupplierCnt = Field(
        default_factory=lambda: defaultdict(int),
        title="不同供应商的真值生产中 frame 数",
    )

    gt_task_success_clip_cnt_by_supplier: GTSupplierCnt = Field(
        default_factory=lambda: defaultdict(int),
        title="不同供应商的真值生产完成 clip 数",
    )
    gt_task_success_frame_cnt_by_supplier: GTSupplierCnt = Field(
        default_factory=lambda: defaultdict(int),
        title="不同供应商的真值生产完成 frame 数",
    )
    gt_valid_rate: float = Field(default=0.0, title="数据有效率")

    def fill_ac_rate(self):
        if self.accept_finished_frame_cnt == 0:
            self.accept_valid_rate = 0.0
        else:
            self.accept_valid_rate = (
                self.accept_valid_frame_cnt / self.accept_finished_frame_cnt
            )

    def fill_gt_rate(self):
        if self.labeled_frame_cnt == 0:
            self.gt_valid_rate = 0.0
        else:
            self.gt_valid_rate = self.gt_task_success_frame_cnt / self.labeled_frame_cnt

    # gt_task_pending_clip_cnt: int = Field(default=0, title="真值生产中 clip 数")
    # gt_task_pending_frame_cnt: int = Field(default=0, title="真值生产中 frame 数")

    # gt_task_finished_clip_cnt: int = Field(default=0, title="真值生产完成 clip 数")
    # gt_task_finished_frame_cnt: int = Field(default=0, title="真值生产完成 frame 数")

    # sensetime_gt_task_finished_clip_cnt: int = Field(
    #     default=0, title="真值生产完成 clip 数(商汤)"
    # )
    # sensetime_gt_task_finished_frame_cnt: int = Field(
    #     default=0, title="真值生产完成 frame 数(商汤)"
    # )

    # imotion_gt_task_finished_frame_cnt: int = Field(
    #     default=0, title="真值生产完成 frame 数(知行)"
    # )
    # imotion_gt_task_finished_frame_cnt: int = Field(
    #     default=0, title="真值生产完成 frame 数(知行)"
    # )

    # zdrive_gt_task_finished_frame_cnt: int = Field(
    #     default=0, title="真值生产完成 frame 数(大卓)"
    # )
    # zdrive_gt_task_finished_frame_cnt: int = Field(
    #     default=0, title="真值生产完成 frame 数(大卓)"
    # )


dashboard_map: dict[int, GTDashboardModel] = {}


def get_dashboard_by_annotate_type(annotate_type: int):
    if annotate_type not in dashboard_map:
        dashboard_map[annotate_type] = GTDashboardModel(
            annotate_type=annotate_type,
            annotate_name=AnnotateTypeEnum.get_chinese_desc_by_code(annotate_type),
        )

    return dashboard_map[annotate_type]


# %%
print("curated frame/clip statistics")
curated_rows = []
with mysql_engine.connect() as conn:
    rows = conn.execute(
        text("""
            select annotate_type, sum(actual_send_label_clip_count) as clip_cnt, sum(actual_send_label_frame_count) as frame_cnt
            from send_label_task_tab
            group by annotate_type
            order by annotate_type
            """)
    )
    for row in rows:
        curated_rows.append(
            {
                "annotate_type": row[0],
                "clip_cnt": int(row[1]),
                "frame_cnt": int(row[2]),
            }
        )

annotate_types = []
for row in curated_rows:
    annotate_types.append(row["annotate_type"])
    cur_item = get_dashboard_by_annotate_type(row["annotate_type"])
    cur_item.curated_clip_cnt += row["clip_cnt"]
    cur_item.curated_frame_cnt += row["frame_cnt"]

# %%
print("labeled frame/clip statistics")
labeled_rows = [
    row
    for row in clip_annotated_record.aggregate(
        [
            {"$match": {"state": {"$in": [2, 4]}}},
            {
                "$group": {
                    "_id": {
                        "annotate_type": "$annotate_type",
                        "is_vis_processed": "$is_vis_processed",
                    },
                    "clip_cnt": {"$sum": 1},
                    "frame_cnt": {"$sum": "$frame_count"},
                }
            },
        ]
    )
]

for row in labeled_rows:
    anno_type, is_vis = (
        row["_id"].get("annotate_type", None),
        row["_id"].get("is_vis_processed", None),
    )
    if anno_type not in annotate_types:
        continue
    dashboard_map[anno_type].labeled_clip_cnt += row["clip_cnt"]
    dashboard_map[anno_type].labeled_frame_cnt += row["frame_cnt"]
    if is_vis is True:
        dashboard_map[anno_type].labeled_vis_processed_clip_cnt += row["clip_cnt"]
        dashboard_map[anno_type].labeled_vis_processed_frame_cnt += row["frame_cnt"]


# %%
print("ac_task pending clip statistics")
ac_task_rows = [
    row
    for row in ac_task_col.aggregate(
        [
            {"$match": {"status": {"$in": [1, 2]}}},
            {
                "$group": {
                    "_id": {
                        "status": "$status",
                        "annotate_type": "$annotate_type",
                    },
                    "clip_cnt": {"$sum": "$clip_ids_num"},
                },
            },
        ]
    )
]

for row in ac_task_rows:
    anno_type, status = row["_id"]["annotate_type"], row["_id"]["status"]
    if anno_type not in annotate_types:
        continue
    dashboard_map[anno_type].accept_pending_clip_cnt += row["clip_cnt"]


# %%
print("ac_task finished clip statistics")
ac_task_clip_rows = [
    row
    for row in check_frame_label_col.aggregate(
        [
            {
                "$group": {
                    "_id": {
                        "annotate_type": "$annotate_type",
                    },
                    "clips": {"$addToSet": "$clip_id"},
                },
            },
            {
                "$project": {
                    "clip_cnt": {"$size": "$clips"},
                }
            },
        ]
    )
]

for row in ac_task_clip_rows:
    anno_type = row["_id"].get("annotate_type", None)
    if anno_type not in annotate_types:
        continue
    dashboard_map[anno_type].accept_finished_clip_cnt += row["clip_cnt"]

# %%
print("ac_task finished frame statistics")
ac_task_frame_standard_rows = [
    row
    for row in check_frame_label_col.aggregate(
        [
            {
                "$group": {
                    "_id": {
                        "annotate_type": "$annotate_type",
                        "standard": "$judging_standard",
                    },
                    "frame_cnt": {"$sum": 1},
                }
            },
        ]
    )
]

for row in ac_task_frame_standard_rows:
    anno_type, standard = (
        row["_id"].get("annotate_type", None),
        row["_id"].get("standard", None),
    )
    if anno_type not in annotate_types or status is None:
        continue
    dashboard_map[anno_type].accept_finished_frame_cnt += row["frame_cnt"]
    if standard in ["A", "B"]:
        dashboard_map[anno_type].accept_valid_frame_cnt += row["frame_cnt"]


# # %%
print("gt task clip statistics")
gt_task_item_rows = [
    row
    for row in gt_task_item_col.aggregate(
        [
            {
                "$group": {
                    "_id": {
                        "annotate_type": "$annotate_type",
                        "status": "$status",
                        "supplier": "$supplier",
                    },
                    "clip_ids": {"$addToSet": "$clip_id"},
                },
            },
            {"$project": {"clip_cnt": {"$size": "$clip_ids"}}},
        ]
    )
]

for row in gt_task_item_rows:
    anno_type, status, supplier = (
        row["_id"].get("annotate_type", None),
        row["_id"].get("status", None),
        row["_id"].get("supplier", None),
    )
    if anno_type not in annotate_types or status is None:
        continue
    if status in ["init", "running"]:
        dashboard_map[anno_type].gt_task_running_clip_cnt += row["clip_cnt"]
        dashboard_map[anno_type].gt_task_running_clip_cnt_by_supplier[supplier] += row[
            "clip_cnt"
        ]
    elif status in ["success"]:
        dashboard_map[anno_type].gt_task_success_clip_cnt += row["clip_cnt"]
        dashboard_map[anno_type].gt_task_success_clip_cnt_by_supplier[supplier] += row[
            "clip_cnt"
        ]


# %%
print("gt task frame statistics")
for row in gt_task_item_col.find(
    {"status": "success"}, {"metadata": 1, "annotate_type": 1, "supplier": 1}
):
    anno_type = row.get("annotate_type", None)
    if anno_type not in annotate_types:
        continue
    if row.get("metadata", None) is None:
        continue
    dashboard_map[anno_type].gt_task_success_frame_cnt += int(
        row.get("metadata", {}).get("NUM_GT", 0)
    )
    dashboard_map[anno_type].gt_task_success_frame_cnt_by_supplier[row["supplier"]] += (
        int(row.get("metadata", {}).get("NUM_GT", 0))
    )


# %% dump to excel

field_titles = {
    name: info.title for name, info in GTDashboardModel.model_fields.items()
}
dashboard_list = list(dashboard_map.values())

for dashboard in dashboard_list:
    dashboard.fill_ac_rate()
    dashboard.fill_gt_rate()


dashboard_json_list = [dashboard.model_dump() for dashboard in dashboard_list]
df = pd.DataFrame(dashboard_json_list)
df.rename(columns=field_titles, inplace=True)

excel_file_path = "gt_dashboard_report.xlsx"
df.to_excel(excel_file_path, index=False)

print(f"Data has been saved to {excel_file_path}")
